"""量化分析 Agent."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Optional

from ..prompts import build_quant_analysis_prompt, load_task_template
from .base import AgentConfig, BaseETFAnalysisAgent


@dataclass
class QuantAnalysisContext:
    fund: str
    current_time: str
    history_summary: str
    feature_summary: str
    risk_summary: str
    quant_result_summary: str | None = None  # 量化分析工具返回的结构化结果


class QuantAnalysisAgent(BaseETFAnalysisAgent):
    """根据量化分析任务提示词生成报告."""

    def __init__(self, *, config: Optional[AgentConfig] = None) -> None:
        super().__init__(config=config)
        self._prompt = build_quant_analysis_prompt()
        self._task_template = load_task_template("quant_analysis_task")

    def build_inputs(self, context: QuantAnalysisContext) -> Dict[str, str]:
        inputs = {
            "current_time": context.current_time,
            "fund": context.fund,
            "history_summary": context.history_summary,
            "feature_summary": context.feature_summary,
            "risk_summary": context.risk_summary,
            "task_description": self._task_template.description,
        }
        # 如果有量化分析结果，添加到输入中
        if context.quant_result_summary:
            inputs["quant_result_summary"] = context.quant_result_summary
        return inputs

    def run(self, context: QuantAnalysisContext) -> str:
        inputs = self.build_inputs(context)
        formatted = self._prompt.format(**inputs)
        return self.llm.invoke(formatted).content  # type: ignore[return-value]

